from ultralytics import YOLO
# 2. 导出ONNX模型
def export_onnx():
    # 加载训练好的最佳模型
    model = YOLO('runs/rail/yolov8n_damage_v1/weights/best.pt')

    # 导出参数配置（严格匹配K230 SDK要求）
    model.export(
        format='onnx',       # 目标格式为ONNX
        opset=12,            # K230支持的OPSET版本（根据SDK文档可能为11/12/13）
        simplify=True,       # 简化模型结构，移除不支持的算子
        dynamic=False,       # 固定输入尺寸（K230不支持动态尺寸）
        imgsz=[640, 640],    # 输入尺寸与训练一致
        batch=1,             # 固定批次大小为1（K230通常单图推理）
        
    )

if __name__ == '__main__':
    export_onnx()        # 训练完成后自动导出ONNX
